#批量处理，脚本，阻塞列队，一个一个跑以免机器吃不消
#首先在input文件夹里面，创建批处理文件夹，把这批量跑的fasta文件全放进去
#输入 python one_by_one.py <文件夹名>即可
#作者：zhuxf

import os
import sys
import numpy
import pickle
import pandas as pd
import time

filePath = sys.argv[1]
#较小的遗传数据库配置 (reduced_dbs) 或完整的遗传数据库配置 (full_dbs)（默认值：' full_dbs '）
db_preset = sys.argv[2]
#filelist = os.listdir("./input/"+filePath)
filelist_temp = os.listdir("./input/"+filePath)
filelist = []

#把奇怪后缀全部换成.fasta
for item in filelist_temp:
    filelist.append(item.split(".")[0] + ".fasta")


#按照数字排序
filelist.sort()
# print(filelist)
filelist3 = []
filelist4 = []
# models = [1, 2, 3, 4, 5]
print("预备执行========================alphafold=====================创建文件夹：single-{output_dir}".format(output_dir=filePath))
#创建批处理的输出文件夹
os.system("mkdir ./output/single-{output_dir}".format(output_dir=filePath))
filelist2 = os.listdir("./output/single-{output_dir}/".format(output_dir=filePath))
for i in filelist2:
    filelist3.append(i+".fasta")
# 差集
for i in filelist:
    if i not in filelist3:
        filelist4.append(i)
    # print(filelist)
    # print(filelist2)
    # print(filelist3)
print("开始执行以下文件：======>")
print(filelist4)

for item in filelist4:
    #print("bash run_alphafold.sh -d ../alphafold_database -o ./output/批处理-{output_dir} -m model_1 -f ./input/{input} -t 2000-01-10".format(output_dir=filePath,input=item))
    print("开始执行===================alphafold===========================文件名："+ item)
#开始跑alphafold
    os.system("bash run_alphafold.sh -d ../alphafold_database -o ./output/single-{output_dir} -f ./input/{dir}/{input} -t 2020-05-14 -c {db_preset}".format(output_dir=filePath,dir=filePath,input=item,db_preset=db_preset))
#修改权限
    os.system("sudo chmod 777 ./output/single-{output_dir}/{output_file}".format(output_dir=filePath,output_file=item.split(".")[0]))
    # filelist3 = []
    # for model in models:
    # #ranked0 的添加plddt
    #     path = './output/批处理-{output_dir}/{output_file}/result_model_{model}.pkl'.format(output_dir=filePath,model=model,output_file=item.split(".")[0])
    #     # 按行读入，删除最后2行
    #     file_old = open('./output/批处理-{output_dir}/{output_file}/ranked_{model}.pdb'.format(output_dir=filePath,model=int(model)-1,output_file=item.split(".")[0]), 'r')
    #     lines = [i for i in file_old]
    #     del lines[-1]
    #     del lines[-1]
    #     file_old.close()
    #     # 再写入临时文件tmp.pdb
    #     file_new = open('./output/批处理-{output_dir}/{output_file}/tmp{model}.pdb'.format(output_dir=filePath,model=model, output_file=item.split(".")[0]),'w')
    #     file_new.write(''.join(lines))
    #     file_new.close()
    #     try:
    #         result_model_1_pkl = open(path, 'rb')
    #     except:
    #         continue
    #     else:
    #         data = pickle.load(result_model_1_pkl)
    #         # 输出plddt参数
    #         data = data["plddt"]
    #         # print(data)
    #         list = data.tolist()
    #         file_tmp = open('./output/批处理-{output_dir}/{output_file}/tmp{model}.pdb'.format(output_dir=filePath,model=model,output_file=item.split(".")[0]), 'r')
    #         tmp_lines = [i for i in file_tmp]
    #         count1 = 0
    #         for index in range(len(tmp_lines)):
    #             if (int(tmp_lines[index][22:26]) == count1 + 1):
    #                 tmp_lines[index] = tmp_lines[index].replace(' 0.00 ', format(float(list[count1]), '.2f') + " ")
    #             elif (int(tmp_lines[index][22:26]) != count1 + 1):
    #                 tmp_lines[index] = tmp_lines[index].replace(' 0.00 ', format(float(list[count1 + 1]), '.2f') + " ")
    #                 count1 = count1 + 1
    #             print(tmp_lines[index])
    #         file_tmp.close()
    #         # print(tmp_lines)
    #
    #         # 再写入临时文件result.pdb
    #         file_result = open('./output/批处理-{output_dir}/{output_file}/ranked_{model}_pLDDT.pdb'.format(output_dir=filePath,model=int(model)-1,output_file=item.split(".")[0]), 'w')
    #         file_result.write(''.join(tmp_lines))
    #         file_result.close()
    os.system("sudo mkdir ./output/single-{output_dir}/{output_file}/{output_file}".format(output_dir=filePath,output_file=item.split(".")[0]))
    os.system("sudo chmod 777 ./output/single-{output_dir}/{output_file}/{output_file}".format(output_dir=filePath,output_file=item.split(".")[0]))
    os.system("sudo cp ./output/single-{output_dir}/{output_file}/rank* ./output/single-{output_dir}/{output_file}/{output_file}".format(output_dir=filePath,output_file=item.split(".")[0]))
    # os.system("sudo rm ./output/批处理-{output_dir}/{output_file}/tmp*".format(output_dir=filePath,output_file=item.split(".")[0]))
